from matplotlib import pyplot as plt
import seaborn as sns
from flow1 import TrafficGenerator
import numpy as np

# 创建TrafficGenerator类的实例
generator = TrafficGenerator(num_nodes=16)

# 调用generate_all_to_all方法生成流量矩阵
traffic_matrix = generator.generate_all_to_all()
# 调用generate_ring_allreduce方法生成流量矩阵
traffic_matrix1 = generator.generate_ring_allreduce()
# 调用generate_brust方法生成流量矩阵
traffic_matrix2 = generator.generate_burst_traffic()
# 打印生成的流量矩阵
# print("生成的All-to-All流量矩阵：")
# print("生成的All-Reduce流量矩阵：")
print("生成的All-Reduce+brust流量矩阵：")
print(traffic_matrix2+traffic_matrix1)

# 将生成的流量矩阵画成热力图
# 设置图片清晰度
# plt.rcParams['figure.dpi'] = 300

# 绘制热力图
# 设置字体大小和图片大小
plt.rcParams['font.size'] = 6  # 设置字体大小为8
plt.figure(figsize=(10, 8))  # 减小图片大小
sns.heatmap(traffic_matrix2+traffic_matrix1, annot=True, fmt=".2f", cmap="YlGnBu",
            xticklabels=generator.tor_nodes, yticklabels=generator.tor_nodes)
# plt.title("All-to-All Traffic Matrix Heatmap")
# plt.title("All-Reduce Traffic Matrix Heatmap")
plt.title("All-Reduce+Brust Traffic Matrix Heatmap")
plt.xlabel("Destination Nodes")
plt.ylabel("Source Nodes")

# 显示图形
plt.show()